import os
import cv2
import numpy as np
import time
from utils import plot_one_box, cal_iou, xyxy_to_xywh, xywh_to_xyxy, updata_trace_list, draw_trace

# 单目标跟踪
# 检测器获得检测框，全程只赋予1个ID，有两个相同的东西进来时，不会丢失唯一跟踪目标
# 检测器的检测框为测量值
# 目标的状态X = [x,y,h,w,delta_x,delta_y],中心坐标，宽高，中心坐标速度
# 观测值
# 如何寻找目标的观测值
# 观测到的是N个框
# 怎么找到目标的观测值
# t时刻的框与t-1后验估计时刻IOU最大的框的那个作为观测值（存在误差，交叉情况下观测值会有误差）
# 所以需要使用先验估计值进行融合
#

# 状态初始化
initial_target_box = [729, 238, 764, 339]  # 目标初始bounding box
# initial_target_box = [193, 342, 250, 474]

initial_box_state = xyxy_to_xywh(initial_target_box)
initial_state = np.array([[initial_box_state[0], initial_box_state[1], initial_box_state[2], initial_box_state[3],
                           0, 0]]).T  # [中心x,中心y,宽w,高h,dx:中心点x坐标变化量(速度),dy:中心点y坐标变化量(速度)]  dx=x-x^-
IOU_Threshold = 0.3  # 匹配时的阈值

# 状态转移矩阵，上一时刻的状态转移到当前时刻...状态转移矩阵需要根据建模确定. 手动列出来方程, 然后获得矩阵.
State_transition_matrix__A = np.array([[1, 0, 0, 0, 1, 0],
                                       [0, 1, 0, 0, 0, 1],
                                       [0, 0, 1, 0, 0, 0],
                                       [0, 0, 0, 1, 0, 0],
                                       [0, 0, 0, 0, 1, 0],
                                       [0, 0, 0, 0, 0, 1]])

# 状态观测矩阵
State_observation_matrix__H = np.eye(6)

# 在设置噪声矩阵的时候，我们通常假设数据的噪声相互之间是独立的，那就只设置对角线上的元素，自己和自己的协方差（也就是“方差”），这时候就只需要考虑数据自己本身的特性了。
# 假设所有噪声都是零均值的白噪声，那么Q和R中的元素就是白噪声的方差
# 如果P0、Q、R无法精确获得，只知道可能的取值范围，则采用可能的较大值（保守）。
# 如果不确切知道Q、R、P0的准确先验信息，应适当增大Q的取值，以增大对实时量测值的利用权重，俗称调谐。但是调谐存在盲目性，无法知道Q要调到多大才行。
# Q和R分别代表对预测值和测量值的置信度（反比），通过影响卡尔曼增益K的值，影响预测值和测量值的权重。越大的R代表越不相信测量值。

# 过程噪声协方差矩阵Q，p(w)~N(0,Process_noise_covariance_matrix__Q)，噪声来自真实世界中的不确定性,
# Q代表建模误差，跟你计算机舍入误差（比如说用了浮点，那么会引入1e-7的误差项）、模型的线性化程度、离散化引入误差有关系；如果你还有输入u的话，输入的误差也会考虑到这里。
# 在跟踪任务当中，过程噪声来自于目标移动的不确定性（突然加速、减速、转弯等）
Process_noise_covariance_matrix__Q = np.eye(6) * 1  # p(w)~N(0,Q)

# 观测噪声协方差矩阵R，p(v)~N(0,Observation_noise_covariance_matrix__R)
# 观测噪声来自于检测框丢失、重叠, 不准确等;
# R的选取与传感器的特性有关，具体到应用中，比如导航系统，我的GPS传感器传下来的位置p会在真值附近有个方差为±3m的噪声，那么R中对应p的那个元素就该是3。
Observation_noise_covariance_matrix__R = np.eye(6) * 0.5

# 控制输入矩阵B
Control_input_matrix__B = None  # 零矩阵
# 状态估计协方差矩阵P初始化  # 随意的设置, 后面自己算
# P的初值就是代表你对你估计的X0有多不确定。P0为0的时候就是说你对你给初值有着坚不可摧的信心。一般来说宁可设成单位阵都不会设成0的。
State_estimation_covariance_matrix__P = np.eye(6)
# State_estimation_covariance_matrix__P = np.random.random((6, 6))*100

if __name__ == "__main__":

    video_path = "./data/testvideo1.mp4"
    label_path = "./data/labels"
    file_name = "testvideo1"
    cap = cv2.VideoCapture(video_path)
    # cv2.namedWindow("track", cv2.WINDOW_NORMAL)
    SAVE_VIDEO = False
    if SAVE_VIDEO:
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        out = cv2.VideoWriter('kalman_output.avi', fourcc, 20, (768, 576))

    # ---------状态初始化----------------------------------------
    frame_counter = 1
    X_posterior = np.array(initial_state)
    P_posterior = np.array(State_estimation_covariance_matrix__P)
    Z = np.array(initial_state)
    trace_list = []  # 用于保存目标box的轨迹

    while (True):
        time.sleep(0.05)
        # Capture frame-by-frame
        ret, frame = cap.read(0.1)

        # last_box_posterior = xywh_to_xyxy(X_posterior[0:4])
        # plot_one_box(last_box_posterior, frame, color=(255, 255, 255), target=False)
        if not ret:
            break
        # print(frame_counter)
        label_file_path = os.path.join(label_path, file_name + "_" + str(frame_counter) + ".txt")
        if os.path.exists(label_file_path):
            with open(label_file_path, "r") as f:
                content = f.readlines()
                max_iou = IOU_Threshold
                max_iou_matched = False  # 记录是否匹配到了目标bbox
                # ---------使用最大IOU来寻找观测值------------
                for j, data_ in enumerate(content):
                    data = data_.replace('\n', "").split(" ")
                    xyxy = np.array(data[1:5], dtype="float")
                    plot_one_box(xyxy, frame)
                    iou = cal_iou(xyxy, xywh_to_xyxy(X_posterior[0:4]))
                    if iou > max_iou:
                        target_box = xyxy
                        max_iou = iou
                        max_iou_matched = True  # 表示匹配到了目标bbox
                if max_iou_matched == True:
                    # 如果找到了最大IOU BOX,则认为该框为观测值
                    plot_one_box(target_box, frame, target=True)
                    xywh = xyxy_to_xywh(target_box)
                    box_center = (int((target_box[0] + target_box[2]) // 2), int((target_box[1] + target_box[3]) // 2))
                    trace_list = updata_trace_list(box_center, trace_list, 80)  # 这里可以动态调节历史轨迹的大小
                    cv2.putText(frame, "Tracking", (int(target_box[0]), int(target_box[1] - 5)), cv2.FONT_HERSHEY_SIMPLEX,
                                0.7,
                                (255, 0, 0), 2)
                    # 计算dx,dy
                    dx = xywh[0] - X_posterior[0]
                    dy = xywh[1] - X_posterior[1]

                    Z[0:4] = np.array([xywh]).T
                    Z[4::] = np.array([dx, dy])

        else:
            print("_-----------------------_")

        if max_iou_matched:
            # -----进行先验估计-----------------
            X_prior = np.dot(State_transition_matrix__A, X_posterior)  # 预测的新状态
            box_prior = xywh_to_xyxy(X_prior[0:4])  # 预测的框
            print("box_prior: %s" % box_prior)
            # plot_one_box(box_prior, frame, color=(0, 0, 0), target=False)
            # -----计算状态估计协方差矩阵P--------
            P_prior_1 = np.dot(State_transition_matrix__A, P_posterior)  # P_posterior: ???
            P_prior = np.dot(P_prior_1, State_transition_matrix__A.T) + Process_noise_covariance_matrix__Q
            # ------计算卡尔曼增益---------------------
            k1 = np.dot(P_prior, State_observation_matrix__H.T)
            k2 = np.dot(np.dot(State_observation_matrix__H, P_prior), State_observation_matrix__H.T) + Observation_noise_covariance_matrix__R
            K = np.dot(k1, np.linalg.inv(k2))
            # --------------后验估计------------
            X_posterior_1 = Z - np.dot(State_observation_matrix__H, X_prior)
            X_posterior = X_prior + np.dot(K, X_posterior_1)
            box_posterior = xywh_to_xyxy(X_posterior[0:4])
            # plot_one_box(box_posterior, frame, color=(255, 255, 255), target=False)
            # ---------更新状态估计协方差矩阵P-----
            P_posterior_1 = np.eye(6) - np.dot(K, State_observation_matrix__H)
            P_posterior = np.dot(P_posterior_1, P_prior)
        else:
            # 如果IOU匹配失败，此时失去观测值，那么直接使用上一次的最优估计作为先验估计
            # 此时直接迭代，不使用卡尔曼滤波
            X_posterior = np.dot(State_transition_matrix__A, X_posterior)
            # X_posterior = np.dot(A_, X_posterior)
            box_posterior = xywh_to_xyxy(X_posterior[0:4])
            # plot_one_box(box_posterior, frame, color=(255, 255, 255), target=False)
            box_center = (
                (int(box_posterior[0] + box_posterior[2]) // 2), int((box_posterior[1] + box_posterior[3]) // 2))
            trace_list = updata_trace_list(box_center, trace_list, 300)
            cv2.putText(frame, "Lost.....................", (box_center[0], box_center[1] - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.7,
                        (255, 0, 0), 2)

        draw_trace(frame, trace_list)  # 仅仅绘制轨迹线

        last_box_posterior = xywh_to_xyxy(X_posterior[0:4])
        plot_one_box(box_posterior, frame, color=(255, 255, 255), target=False)

        # 解释各个框是什么意思.
        cv2.putText(frame, "ALL BOXES(Green)", (25, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 200, 0), 2)
        cv2.putText(frame, "TRACKED BOX(Red)", (25, 75), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        cv2.putText(frame, "Last frame best estimation(White)", (25, 100), cv2.FONT_HERSHEY_SIMPLEX, 0.7,
                    (255, 255, 255), 2)

        cv2.imshow('track', frame)
        if SAVE_VIDEO:
            out.write(frame)
        frame_counter = frame_counter + 1
        if cv2.waitKey(10) & 0xFF == ord('q'):
            break

    # When everything done, release the capture
    time.sleep(10)
    cap.release()
    cv2.destroyAllWindows()

    # 关注我
    # 你关注我了吗
    # 关注一下
